# coding: utf-8
import pandas as pd
from datetime import datetime
import pymongo
import pytz
import os
from time import time

"""
把crawler爬来的历史分钟数据导入mongodb
"""


def _func(df):
    if (df['openPrice'].iloc[0] < 0.0000001):
        return pd.DataFrame([])
    if (df['totalVolume'].sum() < 0.0000001):
        return pd.DataFrame([])
    df['openPrice'].values[1] = df['openPrice'].values[0]
    if df['highPrice'].values[0] > df['highPrice'].values[1]:
        df['highPrice'].values[1] = df['highPrice'].values[0]
    if df['lowPrice'].values[0] < df['lowPrice'].values[1]:
        df['lowPrice'].values[1] = df['lowPrice'].values[0]
    df['totalVolume'].values[1] = df['totalVolume'].values[0] + df['totalVolume'].values[1]
    df['totalValue'].values[1] = df['totalValue'].values[0] + df['totalValue'].values[1]
    return df.iloc[1:]


path = 'D:/wks_python/future_live_data/stock/hdkj/'
filelist = os.listdir(path)
tz = pytz.timezone(pytz.country_timezones('cn')[0])

client = pymongo.MongoClient('121.40.212.219', 27017)  # '121.40.212.219', 27017
db = client.Daily_data
collection = client.Minute_data.one_min_data

#从日志中读取已经写入过的股票跳过。若要全部写入，先清空log.txt。
fi = open('log.txt')
symbols = []
txt = fi.read()
a = txt.replace('\n', ',')
symbols = a.split(',')


fi.close()

print ('begin at %s' % datetime.today().strftime('%Y-%m-%d %H:%M:%S'))
t0 = time()
fi = open('log.txt', 'a')

for f in filelist:
    df = pd.read_csv(path + f)
    if symbols.__contains__(f):
        print("pass" + f)
        continue
    if df.shape[0] == 1 or df.empty:
        continue

    df = df.groupby('dataDate', group_keys=False).apply(_func)
    if not df.empty:
        df = df[
            ['barTime', 'dataDate', 'closePrice', 'openPrice', 'highPrice', 'lowPrice', 'totalVolume', 'totalValue']]
        sec = f.split('.')[0]
        df['symbol'] = sec
        df.rename(columns={'dataDate': 'date', 'closePrice': 'close', \
                           'openPrice': 'open', 'highPrice': 'high', \
                           'lowPrice': 'low', 'totalVolume': 'volume', \
                           'totalValue': 'amount'}, inplace=True)
        df['dateTime'] = df['date'] + ' ' + df['barTime']
        df['dateTime'] = df['dateTime'].apply(lambda x: tz.localize(datetime.strptime(x, '%Y-%m-%d %H:%M')))
        df['volume'] = df['volume'] / 100  # change unit to 手
        df['Fromtonglian'] = True

        df[:1]

        # collection.insert_many(df.to_dict(orient='records'))

        df = df.set_index('dateTime')
        # ---------------write in 5 minute data----------------
        df_day = pd.DataFrame([])
        df_day['open'] = df['open'].resample('D', how='first', closed='right', label='left').dropna()
        df_day['close'] = df['close'].resample('D', how='last', closed='right', label='left').dropna()
        df_day['high'] = df['high'].resample('D', how='max', closed='right', label='left').dropna()
        df_day['low'] = df['low'].resample('D', how='min', closed='right', label='left').dropna()
        df_day['barTime'] = df['barTime'].resample('D', how='last', closed='right', label='left').dropna()
        df_day['date'] = df['date'].resample('D', how='last', closed='right', label='left').dropna()
        df_day['amount'] = df['amount'].resample('D', how='sum', closed='right', label='left').dropna()
        df_day['volume'] = df['volume'].resample('D', how='sum', closed='right', label='left').dropna()
        df_day['symbol'] = sec
        # db.qfq_data.insert_many(df_day.reset_index().to_dict(orient='records'))
        print("5 min done")

        fi.write(f + "\n")

fi.close()
print('finished in %.2fs' % (time() - t0))
